/*
Copyright 2025 PerfXLab (Beijing) Technologies Co., Ltd.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
#include "test_spmv_vnec.hpp"
#define ntest 2000
#define nwarm 1

inline void test_spmv_vnec_fp32(const std::string &input_path, const std::string &result_path, bool iscsc)
{
    Builder<perfSpB_index, perfSpB_index> builder(input_path);
    CSRGraph<perfSpB_index, perfSpB_index> g = builder.MakeGraph();

    perfSpB_index nrows, ncols, nvals;

    nrows = static_cast<perfSpB_index>(g.num_nodes());
    ncols = static_cast<perfSpB_index>(g.num_nodes());
    nvals = static_cast<perfSpB_index>(g.num_edges());

    perfSpB_index *row_indices;
    perfSpB_index *col_indices;
    float *values = new float[nvals];

    row_indices = new perfSpB_index[nvals];
    col_indices = g.out_destValues();
    perfSpB_index *offsets = g.out_offsets();
    perfSpB_index index = 0;
    for (perfSpB_index i = 0; i < nrows; i++)
    {

        for (perfSpB_index j = 0; j < offsets[i + 1] - offsets[i]; j++)
        {
            if (index >= nvals)
            {
                std::cerr << "Error index >= nvals" << index << " " << nvals << std::endl;
                exit(-1);
            }
            values[index] = i;
            row_indices[index++] = static_cast<perfSpB_index>(i);
        }
    }

    PerfSpB_init();
    perfSpB_object_format_ext format;
    format.matrix_format = perfSpB_csr;
    PerfSpB_Global_Option_set_ext(perfSpB_matrix_format, format);
    perfSpB_matrix G = NULL;
    perfSpB_info info;
    SAFE_CALL(perfSpB_matrix_new(&G, perfSpB_FP32, nrows, ncols));
    perfSpB_matrix_build_FP32(G, row_indices, col_indices, values, nvals);
    perfSpB_vector x;
    SAFE_CALL(perfSpB_vector_new(&x, perfSpB_FP32, ncols));
    for (perfSpB_index j = 0; j < ncols; j++)
    {
        perfSpB_vector_set_element_FP32(x, j, j);
    }
    perfSpB_vector result_benchmark;
    SAFE_CALL(perfSpB_vector_new(&result_benchmark, perfSpB_FP32, nrows));
    perfSpB_vector result_opt;
    SAFE_CALL(perfSpB_vector_new(&result_opt, perfSpB_FP32, nrows));
    struct timeval opt_start, opt_end;
    for (perfSpB_index i = 0; i < 1; i++)
    {
        PerfSpB_mxv_FP32_serial(result_benchmark, G, x);
    }
    //---------------------------------------------------------------------------------------------------------------------------
    if (nrows < 3000000)
    {
        PerfSpB_VNEC_type vnec_t = IRD_VNEC(G, perfSpB_FP32);
        if (vnec_t == PerfSpB_VNEC_D)
        {
            printf("-D is selected\n");
            VNEC_D_FP32 *mat_thd = Build_VNEC_D_FP32(G);
            // warm up
            for (perfSpB_index i = 0; i < nwarm; i++)
            {
                SpMV_VNEC_D_FP32(result_opt, G, x, mat_thd);
            }
            gettimeofday(&opt_start, NULL);
            for (perfSpB_index i = 0; i < ntest; i++)
            {
                SpMV_VNEC_D_FP32(result_opt, G, x, mat_thd);
            }
            gettimeofday(&opt_end, NULL);
            VNEC_D_free_fp32(mat_thd);
        }
        else
        {
            printf("-S is selected\n");
            VNEC_S_FP32 *mat_thd = Build_VNEC_S_FP32(G);
            // warm up
            for (perfSpB_index i = 0; i < nwarm; i++)
            {
                SpMV_VNEC_S_FP32(result_opt, G, x, mat_thd);
            }
            gettimeofday(&opt_start, NULL);
            for (perfSpB_index i = 0; i < ntest; i++)
            {
                SpMV_VNEC_S_FP32(result_opt, G, x, mat_thd);
            }
            gettimeofday(&opt_end, NULL);
            VNEC_S_free_fp32(mat_thd);
        }
    }
    else
    {
        printf("-L is selected\n");
        VNEC_L *mat_thd = Build_VNEC_L(G);
        // warm up
        for (perfSpB_index i = 0; i < nwarm; i++)
        {
            SpMV_VNEC_L_FP32(result_opt, G, x, mat_thd);
        }
        gettimeofday(&opt_start, NULL);
        for (perfSpB_index i = 0; i < ntest; i++)
        {
            SpMV_VNEC_L_FP32(result_opt, G, x, mat_thd);
        }
        gettimeofday(&opt_end, NULL);
        free(mat_thd->diagonal_start);
        free(mat_thd->diagonal_end);
        free(mat_thd->nz_indices);
        free(mat_thd->thread_coord_start);
        free(mat_thd->thread_coord_end);
        free(mat_thd);
    }
    long opt_timeuse = 1000000 * (opt_end.tv_sec - opt_start.tv_sec) + opt_end.tv_usec - opt_start.tv_usec;
    printf("\n ------------------------------------------------------------------------------");
    printf("\n SpMV overall time elapsed: %f ms\n", opt_timeuse / 1000.0 / ntest);
    printf(" ------------------------------------------------------------------------------ \n");
    float diff = 0.0;
    for (perfSpB_index i = 0; i < nrows; ++i)
    {
        float value_benchmark, value_opt;
        perfSpB_vector_extract_element_FP32(&value_benchmark, result_benchmark, i);
        perfSpB_vector_extract_element_FP32(&value_opt, result_opt, i);
        diff = fabs(value_benchmark - value_opt);
        ASSERT_LE(diff, value_benchmark * 1e-4);
    }

    perfSpB_matrix_free(&G);
    perfSpB_vector_free(&x);
    perfSpB_vector_free(&result_opt);
    perfSpB_vector_free(&result_benchmark);
}

inline void test_spmv_vnec_fp64(const std::string &input_path, const std::string &result_path, bool iscsc)
{
    Builder<perfSpB_index, perfSpB_index> builder(input_path);
    CSRGraph<perfSpB_index, perfSpB_index> g = builder.MakeGraph();

    perfSpB_index nrows, ncols, nvals;

    nrows = static_cast<perfSpB_index>(g.num_nodes());
    ncols = static_cast<perfSpB_index>(g.num_nodes());
    nvals = static_cast<perfSpB_index>(g.num_edges());

    perfSpB_index *row_indices;
    perfSpB_index *col_indices;
    double *values = new double[nvals];

    row_indices = new perfSpB_index[nvals];
    col_indices = g.out_destValues();
    perfSpB_index *offsets = g.out_offsets();
    perfSpB_index index = 0;
    for (perfSpB_index i = 0; i < nrows; i++)
    {

        for (perfSpB_index j = 0; j < offsets[i + 1] - offsets[i]; j++)
        {
            if (index >= nvals)
            {
                std::cerr << "Error index >= nvals" << index << " " << nvals << std::endl;
                exit(-1);
            }
            values[index] = i;
            row_indices[index++] = static_cast<perfSpB_index>(i);
        }
    }
    PerfSpB_init();
    perfSpB_object_format_ext format;
    format.matrix_format = perfSpB_csr;
    PerfSpB_Global_Option_set_ext(perfSpB_matrix_format, format);
    perfSpB_matrix G = NULL;
    perfSpB_info info;
    SAFE_CALL(perfSpB_matrix_new(&G, perfSpB_FP64, nrows, ncols));
    perfSpB_matrix_build_FP64(G, row_indices, col_indices, values, nvals);
    perfSpB_vector x;
    SAFE_CALL(perfSpB_vector_new(&x, perfSpB_FP64, ncols));
    for (perfSpB_index j = 0; j < ncols; j++)
    {
        perfSpB_vector_set_element_FP64(x, j, j);
    }
    perfSpB_vector result_benchmark;
    SAFE_CALL(perfSpB_vector_new(&result_benchmark, perfSpB_FP64, nrows));
    perfSpB_vector result_opt;
    SAFE_CALL(perfSpB_vector_new(&result_opt, perfSpB_FP64, nrows));
    struct timeval opt_start, opt_end;
    for (perfSpB_index i = 0; i < 1; i++)
    {
        PerfSpB_mxv_FP64_serial(result_benchmark, G, x);
    }
    //---------------------------------------------------------------------------------------------------------------------------
    if (nrows < 3000000)
    {
        PerfSpB_VNEC_type vnec_t = IRD_VNEC(G, perfSpB_FP64);
        if (vnec_t == PerfSpB_VNEC_D)
        {
            printf("-D is selected\n");
            VNEC_D_FP64 *mat_thd = Build_VNEC_D_FP64(G);
            // warm up
            for (perfSpB_index i = 0; i < nwarm; i++)
            {
                SpMV_VNEC_D_FP64(result_opt, G, x, mat_thd);
            }
            gettimeofday(&opt_start, NULL);
            for (perfSpB_index i = 0; i < ntest; i++)
            {
                SpMV_VNEC_D_FP64(result_opt, G, x, mat_thd);
            }
            gettimeofday(&opt_end, NULL);
            VNEC_D_free_fp64(mat_thd);
        }
        else
        {
            printf("-S is selected\n");
            VNEC_S_FP64 *mat_thd = Build_VNEC_S_FP64(G);
            // warm up
            for (perfSpB_index i = 0; i < nwarm; i++)
            {
                SpMV_VNEC_S_FP64(result_opt, G, x, mat_thd);
            }
            gettimeofday(&opt_start, NULL);
            for (perfSpB_index i = 0; i < ntest; i++)
            {
                SpMV_VNEC_S_FP64(result_opt, G, x, mat_thd);
            }
            gettimeofday(&opt_end, NULL);
            VNEC_S_free_fp64(mat_thd);
        }
    }
    else
    {
        printf("-L is selected\n");
        VNEC_L *mat_thd = Build_VNEC_L(G);
        // warm up
        for (perfSpB_index i = 0; i < nwarm; i++)
        {
            SpMV_VNEC_L_FP64(result_opt, G, x, mat_thd);
        }
        gettimeofday(&opt_start, NULL);
        for (perfSpB_index i = 0; i < ntest; i++)
        {
            SpMV_VNEC_L_FP64(result_opt, G, x, mat_thd);
        }
        gettimeofday(&opt_end, NULL);
        free(mat_thd->diagonal_start);
        free(mat_thd->diagonal_end);
        free(mat_thd->nz_indices);
        free(mat_thd->thread_coord_start);
        free(mat_thd->thread_coord_end);
        free(mat_thd);
    }
    long opt_timeuse = 1000000 * (opt_end.tv_sec - opt_start.tv_sec) + opt_end.tv_usec - opt_start.tv_usec;
    printf("\n ------------------------------------------------------------------------------");
    printf("\n VNEC_FP64 mxv overall time elapsed: %f ms\n", opt_timeuse / 1000.0 / ntest);
    printf(" ------------------------------------------------------------------------------ \n");
    double diff = 0.0;
    for (perfSpB_index i = 0; i < nrows; ++i)
    {
        double value_benchmark, value_opt;
        perfSpB_vector_extract_element_FP64(&value_benchmark, result_benchmark, i);
        perfSpB_vector_extract_element_FP64(&value_opt, result_opt, i);
        diff = fabs(value_benchmark - value_opt);
        ASSERT_LE(diff, value_benchmark * 1e-8);
    }

    perfSpB_matrix_free(&G);
    perfSpB_vector_free(&x);
    perfSpB_vector_free(&result_opt);
    perfSpB_vector_free(&result_benchmark);
}

inline void test_read_fp32(const std::string &sg_path, const std::string &m_path)
{
    Builder<perfSpB_index, perfSpB_index> builder(sg_path);
    CSRGraph<perfSpB_index, perfSpB_index> g = builder.MakeGraph();
    perfSpB_index nrows, ncols, nvals;
    nrows = static_cast<perfSpB_index>(g.num_nodes());
    ncols = static_cast<perfSpB_index>(g.num_nodes());
    nvals = static_cast<perfSpB_index>(g.num_edges());

    perfSpB_index *row_indices;
    perfSpB_index *col_indices;
    float *values = new float[nvals];

    row_indices = new perfSpB_index[nvals];
    col_indices = g.out_destValues();
    perfSpB_index *offsets = g.out_offsets();
    perfSpB_index index = 0;
    for (perfSpB_index i = 0; i < nrows; i++)
    {

        for (perfSpB_index j = 0; j < offsets[i + 1] - offsets[i]; j++)
        {
            if (index >= nvals)
            {
                std::cerr << "Error index >= nvals" << index << " " << nvals << std::endl;
                exit(-1);
            }
            values[index] = 1;
            row_indices[index++] = static_cast<perfSpB_index>(i);
        }
    }

    PerfSpB_init();
    perfSpB_object_format_ext format;
    format.matrix_format = perfSpB_csr;
    PerfSpB_Global_Option_set_ext(perfSpB_matrix_format, format);
    perfSpB_matrix G = NULL;
    perfSpB_info info;
    SAFE_CALL(perfSpB_matrix_new(&G, perfSpB_FP32, nrows, ncols));
    perfSpB_matrix_build_FP32(G, row_indices, col_indices, values, nvals);

    std::vector<perfSpB_index> row_indices_;
    std::vector<perfSpB_index> col_indices_;
    std::vector<float> values_;
    perfSpB_index nrows_, ncols_, nvals_;
    const char *mtx_path = m_path.c_str();

    readMtx_coo<float>(mtx_path, &row_indices_, &col_indices_, &values_, &nrows_, &ncols_, &nvals_, 1, false);

    perfSpB_matrix G_ = NULL;
    SAFE_CALL(perfSpB_matrix_new(&G_, perfSpB_FP32, nrows_, ncols_));
    perfSpB_matrix_build_FP32(G_, &row_indices_[0], &col_indices_[0], (float *)&values_[0], nvals_);

    perfSpB_vector x;
    SAFE_CALL(perfSpB_vector_new(&x, perfSpB_FP32, ncols));
    for (perfSpB_index j = 0; j < ncols; j++)
    {
        perfSpB_vector_set_element_FP32(x, 1, j);
    }
    perfSpB_vector result_mm;
    SAFE_CALL(perfSpB_vector_new(&result_mm, perfSpB_FP32, nrows));
    perfSpB_vector result_sg;
    SAFE_CALL(perfSpB_vector_new(&result_sg, perfSpB_FP32, nrows));

    PerfSpB_mxv_FP32_serial(result_sg, G, x);
    PerfSpB_mxv_FP32_serial(result_mm, G_, x);
    float diff = 0.0;
    float sum_mm = 0.0;
    float sum_sg = 0.0;
    for (perfSpB_index i = 1; i < nrows; ++i)
    {
        float value_sg, value_mm;
        perfSpB_vector_extract_element_FP32(&value_sg, result_sg, i);
        perfSpB_vector_extract_element_FP32(&value_mm, result_mm, i);
        // diff = fabs(value_sg - value_mm);
        // ASSERT_LE(diff, value_mm * 1e-4);
        sum_mm += value_mm;
        sum_sg += value_sg;
        if(value_mm != value_sg)
        {
            printf("\n (%ld)  mm=%f sg = %f  \n",i,value_mm,value_sg);
        }
        ASSERT_EQ(value_mm, value_sg);
    }
    // printf("\n mm=%f sg = %f  \n",sum_mm,sum_sg);

    perfSpB_matrix_free(&G);
    perfSpB_vector_free(&x);
    perfSpB_vector_free(&result_sg);
    perfSpB_vector_free(&result_mm);

    // float diff = 0.0;
    // float sum_mm = 0.0;
    // float sum_sg = 0.0;
    // for (perfSpB_index i = 0; i < 1; ++i)
    // {
    //     for (perfSpB_index j = 0; j < 1; ++j)
    //     {
    //         float value_mm, value_sg;
    //         perfSpB_matrix_extract_element_FP32(&value_mm, G_, i, j);
    //         perfSpB_matrix_extract_element_FP32(&value_sg, G, i, j);
    //         sum_mm += value_mm;
    //         sum_sg += value_sg;
    //         // diff = fabs(value_mm - value_sg);
    //         // ASSERT_LE(diff, value_mm * 1e-4);
    //         // if(value_mm != value_sg)
    //         // {
    //         //     printf("\n (%ld,%ld)  mm=%f sg = %f  \n",i,j,value_mm,value_sg);
    //         // }
    //     }
    // }
    // printf("\n mm=%f sg = %f  \n",sum_mm,sum_sg);

    // perfSpB_matrix_free(&G);
    perfSpB_matrix_free(&G_);

}
#define BATCH_TEST(case_, name)	\
TEST(mxv_test_VNEC_FP32, case_)	\
{	\
    std::string input_file(MATRIX_SG_FULL_PATH(name));	\
    std::string output_file(RESULT_FULL_PATH(name));	\
    test_spmv_vnec_fp32(input_file, output_file, false);\
}	\
TEST(mxv_test_VNEC_FP64, case_)	\
{	\
    std::string input_file(MATRIX_SG_FULL_PATH(name));	\
    std::string output_file(RESULT_FULL_PATH(name));	\
    test_spmv_vnec_fp64(input_file, output_file, false);\
}	
BATCH_TEST(amazon0302, amazon0302)
BATCH_TEST(ca_CondMat, ca-CondMat)
BATCH_TEST(email_EuAll, email-EuAll)
//BATCH_TEST(Oregon1, Oregon-1)
//BATCH_TEST(Oregon2, Oregon-2)
//BATCH_TEST(p2pGnutella04, p2p-Gnutella04)
//BATCH_TEST(p2pGnutella05, p2p-Gnutella05)
//BATCH_TEST(p2pGnutella06, p2p-Gnutella06)
//BATCH_TEST(p2pGnutella08, p2p-Gnutella08)
//BATCH_TEST(p2pGnutella09, p2p-Gnutella09)
//BATCH_TEST(p2pGnutella24, p2p-Gnutella24)
//BATCH_TEST(p2pGnutella25, p2p-Gnutella25)
//BATCH_TEST(p2pGnutella30, p2p-Gnutella30)
//BATCH_TEST(p2pGnutella31, p2p-Gnutella31)
BATCH_TEST(roadNetPA, roadNet-PA)
BATCH_TEST(roadNetTX, roadNet-TX)
BATCH_TEST(socEpinions1, soc-Epinions1)
//BATCH_TEST(socLiveJournal1, soc-LiveJournal1) //slow
BATCH_TEST(socsignepinions, soc-sign-epinions)
BATCH_TEST(socsignSlashdot081106, soc-sign-Slashdot081106)
BATCH_TEST(socsignSlashdot090216, soc-sign-Slashdot090216)
BATCH_TEST(socsignSlashdot090221, soc-sign-Slashdot090221)
BATCH_TEST(socSlashdot0811, soc-Slashdot0811)
//BATCH_TEST(socSlashdot0902, soc-Slashdot0902) //segement failed
BATCH_TEST(webNotreDame, web-NotreDame)
//BATCH_TEST(wikiTalk, wiki-Talk)
//BATCH_TEST(wikiVote, wiki-Vote)
